from matplotlib import pyplot as plt
from torchvision import transforms, utils
import torch
from tqdm import tqdm
from densenet import densenet121
import pickle
import sys
sys.path.append("/home/hzh/.local/lib/python3.5/site-packages")
import cv2


class Predictor:

    def __init__(self, net):
        self._net = net
        with open("./data/chi3500.txt", 'r') as f:
            self.chi3500 = f.readline()

        with open('idx2class.pickle', 'rb') as handle:
            self.idx2class = pickle.load(handle)

    def predict(self, img):
        assert img.shape == (50, 50, 3), "wrong shape of img:%s"%img.shape
        img = img.transpose((2, 0, 1))
        img = torch.from_numpy(img).reshape(1, 3, 50, 50)
        # img_np = img.numpy().reshape(3, 50, 50).transpose(1, 2, 0)
        # cv2.imshow("test", img_np)
        # cv2.waitKey(0)
        outputs = self._net(img)
        _, predicted = torch.max(outputs.data, 1)
        print(predicted)
        return self.label2chi(predicted.cpu())

    def label2chi(self, idx):
        return self.chi3500[int(self.idx2class[idx.item()])]



def load_pytorch_dense_net():
    net = densenet121(num_classes=3800)
    net.load_state_dict(torch.load("./dense_net_param.pth", map_location=torch.device('cpu')))
    net.eval()
    return Predictor(net)
